import numpy as np
import h5py
import os

xs = []
ys = []
for name in sorted(os.listdir('out/feature')):
  # xxxxyyzz.h5
  # 12345678901
  if len(name) != 11 or not name.startswith('2020'):
    continue
  with h5py.File(f'out/feature/{name}', 'r') as f:
    xs.append(f['x'][...])
    ys.append(f['y'][...])

x = np.concatenate(xs, axis=0)
y = np.concatenate(ys, axis=0)
assert len(x) == len(y)
assert len(x) > 0
n = len(x)

x_train, x_val, x_test = np.split(x, [int(n * 0.8), int(n * 0.9)])
y_train, y_val, y_test = np.split(y, [int(n * 0.8), int(n * 0.9)])

with h5py.File(f'out/feature/train.h5', 'w') as f:
  f['x'] = x_train
  f['y'] = y_train

with h5py.File(f'out/feature/val.h5', 'w') as f:
  f['x'] = x_val
  f['y'] = y_val

with h5py.File(f'out/feature/test.h5', 'w') as f:
  f['x'] = x_test
  f['y'] = y_test
